import os
import os.path as osp
import platform
from time import time
import argparse
from glob import glob
from tqdm import tqdm
import numpy as np
from PIL import Image, ImageDraw
import tflite_runtime.interpreter as tflite
import torch
import torchvision
import torchvision.transforms as transforms
from utils import *


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--model', type=str, default=osp.join('data', 'mnist_net_edgetpu.tflite'))
    parser.add_argument('--dataset', type=str, choices=['mnist', 'cifar10'])
    args = parser.parse_args()

    interpreter = tflite.Interpreter(
        model_path=args.model)
    interpreter.allocate_tensors()

    _, test_loader = get_test_dataset(args.dataset)

    for i, (image, label) in enumerate(test_loader, 1):
        image = (image[0].numpy() * 255).astype(np.uint8)
        set_input(
            interpreter=interpreter,
            image=image)
        interpreter.invoke()
        y = get_output(
            interpreter=interpreter)
        print('label={}, y={}'.format(label, np.argmax(y)))


if __name__ == '__main__':
    main()
